// Date:   Thu Aug 14 03:13:30 PM 2025
// Mail:   lunar_ubuntu@qq.com
// Author: https://github.com/xiaoqixian

#include "sampler.hpp"
#include "av.hpp"
#include "common.hpp"
#include "datapipe.hpp"
#include <atomic>
#include <cassert>
#include <cstdint>
#include <opencv2/core/operations.hpp>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

extern "C" {
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/opt.h>
#include <libavutil/time.h>
#include <libswscale/swscale.h>
}

#define WAIT_INTERVAL 1
#define RECORD_PLAYTIME 1

[[maybe_unused]]
static std::chrono::milliseconds calc_pts_diff_ms(int64_t pts1, int64_t pts2, const AVStream* stream) {
  int64_t delta_pts = pts2 - pts1;
  double delta_seconds = delta_pts * av_q2d(stream->time_base);

  return std::chrono::milliseconds(
    static_cast<int64_t>(delta_seconds * 1000.0)
  );
}


VideoStream::VideoStream(Config config): config_(config) {
  url_ = config_["url"].as<std::string>();
  label_ = config_["label"].as<int>();
  src_idx_ = config_["src_idx"].as<int>();

  if (avformat_open_input(&fmt_ctx_, url_.c_str(), nullptr, nullptr) < 0) {
    throw std::runtime_error("Failed to open input video");
  }
  if (avformat_find_stream_info(fmt_ctx_, nullptr) < 0) {
    throw std::runtime_error("Failed to find stream info");
  }
  // Find the first video stream
  for (size_t i = 0; i < fmt_ctx_->nb_streams; i++) {
    if (fmt_ctx_->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
      video_stream_index_ = i;
      break;
    }
  }
  if (video_stream_index_ < 0) throw std::runtime_error("No video stream found");

  // Initialize decoder
  codec_par_ = fmt_ctx_->streams[video_stream_index_]->codecpar;
  codec_ = avcodec_find_decoder(codec_par_->codec_id);
  if (codec_ == nullptr) {
    throw std::runtime_error("Failed to find a decoder");
  }
}

VideoStream::~VideoStream() {
  if (sws_ctx_) {
    sws_freeContext(sws_ctx_);
    sws_ctx_ = nullptr;
  }
  if (fmt_ctx_) {
    avformat_close_input(&fmt_ctx_);
  }
}

VideoStream::VideoStream(VideoStream&& other) noexcept
  : config_(std::move(other.config_)),
    fmt_ctx_(std::exchange(other.fmt_ctx_, nullptr)),
    codec_par_(std::exchange(other.codec_par_, nullptr)),
    codec_(std::exchange(other.codec_, nullptr)),
    codec_ctx_map_(std::move(other.codec_ctx_map_)),
    video_stream_index_(std::exchange(other.video_stream_index_, -1))
{}

VideoStream& VideoStream::operator=(VideoStream&& other) noexcept {
  if (this != &other) {
    if (fmt_ctx_) avformat_close_input(&fmt_ctx_);

    config_ = std::move(other.config_);
    fmt_ctx_ = std::exchange(other.fmt_ctx_, nullptr);
    codec_ctx_map_ = std::move(other.codec_ctx_map_);
    codec_par_ = std::exchange(other.codec_par_, nullptr);
    codec_ = std::exchange(other.codec_, nullptr);
    video_stream_index_ = std::exchange(other.video_stream_index_, -1);
  }
  return *this;
}

void VideoStream::rewind() {
  av_seek_frame(fmt_ctx_, video_stream_index_, 0, AVSEEK_FLAG_BACKWARD);
}

AVDecoder::AVDecoder(int id, Config srcs_cfg, std::atomic_bool* stopped, DataPipe<DecodeRequest>* req_ch)
  : decoder_id_(id)
  , stopped_(stopped)
  , req_ch_(req_ch)
{
  codec_ctx_list_.resize(srcs_cfg.size());
  av_frame_ = av_frame_alloc();

  for (auto const src: srcs_cfg) {
    AVFormatContext* fmt_ctx = nullptr;
    int video_stream_index = -1;
    auto const url = src["url"].as<std::string>();
    auto const src_idx = src["src_idx"].as<int>();

    if (avformat_open_input(&fmt_ctx, url.c_str(), nullptr, nullptr) < 0) {
      throw std::runtime_error("Failed to open input video");
    }
    if (avformat_find_stream_info(fmt_ctx, nullptr) < 0) {
      throw std::runtime_error("Failed to find stream info");
    }
    // Find the first video stream
    for (size_t i = 0; i < fmt_ctx->nb_streams; i++) {
      if (fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
        video_stream_index = i;
        break;
      }
    }
    if (video_stream_index < 0) throw std::runtime_error(cv::format("No video stream found in url: %s", url.c_str()));

    // Initialize decoder
    auto* codec_par = fmt_ctx->streams[video_stream_index]->codecpar;
    auto* codec = avcodec_find_decoder(codec_par->codec_id);
    if (codec == nullptr) {
      throw std::runtime_error("Failed to find a decoder");
    }
    
    codec_ctx_list_[src_idx] = AVCodecCtxWrapper(codec, codec_par);
  }

  worker_th_ = std::thread(AVDecoder::worker, this, decoder_id_);
}

#include <utility>  // for std::exchange

AVDecoder::AVDecoder(AVDecoder&& other) noexcept
  : codec_ctx_list_(std::move(other.codec_ctx_list_)),
    stopped_(std::exchange(other.stopped_, nullptr)),
    worker_th_(std::move(other.worker_th_)),
    req_ch_(std::exchange(other.req_ch_, nullptr)),
    av_frame_(std::exchange(other.av_frame_, nullptr))
{}

AVDecoder& AVDecoder::operator=(AVDecoder&& other) noexcept {
  if (this != &other) {
    if (worker_th_.joinable()) {
      worker_th_.join();
    }

    codec_ctx_list_ = std::move(other.codec_ctx_list_);
    stopped_   = std::exchange(other.stopped_, nullptr);
    worker_th_ = std::move(other.worker_th_);
    req_ch_    = std::exchange(other.req_ch_, nullptr);
    av_frame_  = std::exchange(other.av_frame_, nullptr);
  }
  return *this;
}

AVDecoder::~AVDecoder() {
  if (av_frame_ != nullptr) {
    av_frame_free(&av_frame_);
  }
}

// single thread only
// if fail, the returned cv_mat is empty
Frame AVDecoder::decode_keyframe(FramePkts const& pkts) {
  Frame res = static_cast<FrameBase const&>(pkts);
  auto* codec_ctx = codec_ctx_list_[pkts.src_idx].codec_ctx_;
  codec_ctx->skip_frame = AVDiscard::AVDISCARD_NONKEY;

  auto guard = make_scope_guard([codec_ctx]() {
    avcodec_flush_buffers(codec_ctx);
    codec_ctx->skip_frame = AVDiscard::AVDISCARD_DEFAULT;
  });

  for (size_t i = 0; i < pkts.inter_pkts.size(); i++) {
    const int ret = avcodec_send_packet(codec_ctx, &pkts.inter_pkts[i].pkt_);
    if (ret != 0) {
      throw_av_error(ret);
    }

    while (true) {
      const int ret2 = avcodec_receive_frame(codec_ctx, av_frame_);
      if (ret2 >= 0) {
        res.cv_mat = avframe_to_cvmat(av_frame_);
        return res;
      }
      else if (ret2 == -EAGAIN) {
        break;
      }
      else {
        throw_av_error(ret2);
      }
    }
  }
  return res;
}

std::vector<Frame> AVDecoder::decode_frames(FramePkts const& pkts) {
  std::vector<Frame> res;
  auto* codec_ctx = codec_ctx_list_[pkts.src_idx].codec_ctx_;
  for (size_t i = 0; i < pkts.inter_pkts.size(); i++) {
    const int ret = avcodec_send_packet(codec_ctx, &pkts.inter_pkts[i].pkt_);
    if (ret != 0) {
      throw_av_error(ret);
    }

    while (true) {
      const int ret2 = avcodec_receive_frame(codec_ctx, av_frame_);
      if (ret2 >= 0) {
        Frame frame = static_cast<FrameBase const&>(pkts);
        frame.cv_mat = avframe_to_cvmat(av_frame_);
        res.emplace_back(std::move(frame));
      }
      else if (ret2 == -EAGAIN) {
        break;
      }
      else {
        throw_av_error(ret2);
      }
    }
  }
  return res;
}

void AVDecoder::worker(AVDecoder* self, const int id) {
  auto& req_ch = *(self->req_ch_);
  std::vector<Frame> buf;
  while (!self->stopped_->load(std::memory_order_relaxed)) {
    auto req_opt = req_ch.block_pop_front_unless(*self->stopped_);
    if (!req_opt.has_value()) {
      assert(self->stopped_->load(std::memory_order_relaxed));
      break;
    }
    
    auto& req = req_opt.value();
    switch (req.req_type) {
      case DecodeRequest::ReqType::KeyFrame: {
        auto frame = self->decode_keyframe(*req.pkts);
        buf.emplace_back(std::move(frame));
        break;
      }
      case DecodeRequest::ReqType::Frames: {
        buf = self->decode_frames(*req.pkts);
        break;
      }
    }
    req.res_promise.set_value(std::move(buf));
  }

  printf("Decoder %d exit\n", id);
}

ChannelWorker::ChannelWorker(
  int id,
  Config ch_cfg,
  std::unordered_map<std::string, Config> const& srcs_cfg,
  DataPipe<FramePkts>& window,
  std::unordered_map<int, std::atomic_size_t>& total_emit_cnt
)
  : config_(ch_cfg)
  , channel_id_(id)
  , window_(window)
  , total_emit_cnt_(total_emit_cnt)
{
  auto const srcs_num = config_["srcs"].size();
  srcs_.reserve(srcs_num);

  for (auto const& src_id: config_["srcs"]) {
    auto const src_id_str = src_id.as<std::string>();
    srcs_.emplace_back(srcs_cfg.find(src_id_str)->second);
  }
}

ChannelWorker::~ChannelWorker() {
  for (auto const [label, cnt]: emit_seq_) {
    total_emit_cnt_[label].fetch_add(cnt, std::memory_order_relaxed);
  }
}

void ChannelWorker::start() {
  int src_idx = yaml_get_or_default(config_, "start_idx", 0);
  const int loop_times = yaml_get_or_default(config_, "loop_times", 1) * srcs_.size();

  auto* stream = &srcs_[src_idx];
  FramePkts frame_pkts;

#if WAIT_INTERVAL
  auto last_pkt_pts = AV_NOPTS_VALUE;
  auto last_emplace_ts = std::chrono::high_resolution_clock::now();
#endif

#if RECORD_PLAYTIME
  auto video_starttime = std::chrono::high_resolution_clock::now();
#endif

  for (int lp = 0; lp < loop_times; lp++) {
    printf("Channel %d playing %s\n", channel_id_, stream->url_.c_str());
    int read_frame_ret = 0;
    while (true) {
      AVPacketWrapper pkt_wrapper;
      AVPacket* av_pkt = &pkt_wrapper.pkt_;

      read_frame_ret = av_read_frame(stream->fmt_ctx_, av_pkt);
      if (read_frame_ret < 0) {
        break;
      }
      if (av_pkt->stream_index != stream->video_stream_index_) continue;

      if (av_pkt->flags & AV_PKT_FLAG_KEY) {
        if (!frame_pkts.inter_pkts.empty()) {
          assert(frame_pkts.channel_id != -1);
          window_.emplace_back(std::move(frame_pkts));
        }

        frame_pkts.channel_id = channel_id_;
        frame_pkts.src_idx = stream->src_idx_;
        frame_pkts.ms_ts = now_ms();
        frame_pkts.label = stream->label_;
        frame_pkts.emit_seq = emit_seq_[stream->label_];

#if WAIT_INTERVAL
        if (last_pkt_pts != AV_NOPTS_VALUE) {
          auto const key_pkt_diff_ms = calc_pts_diff_ms(last_pkt_pts, av_pkt->pts, 
            stream->fmt_ctx_->streams[stream->video_stream_index_]);
          auto const emplace_diff_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::high_resolution_clock::now() - last_emplace_ts
          );
          auto const diff = key_pkt_diff_ms - emplace_diff_ms;
          if (diff > std::chrono::milliseconds(0)) {
            std::this_thread::sleep_for(diff);
          }
        }
        
        last_pkt_pts = av_pkt->pts;
        last_emplace_ts = std::chrono::high_resolution_clock::now();
#endif
      }

      frame_pkts.inter_pkts.emplace_back(std::move(pkt_wrapper));
    }

    if (read_frame_ret < 0 && read_frame_ret != AVERROR_EOF) {
      throw_av_error(read_frame_ret);
    }
    else if (read_frame_ret == AVERROR_EOF) {
      stream->rewind();
    }

#if RECORD_PLAYTIME
    auto const video_playtime = std::chrono::duration_cast<std::chrono::seconds>(
      std::chrono::high_resolution_clock::now() - video_starttime
    );
    printf("Channel %d: video %s playtime: %zds\n", channel_id_, stream->url_.c_str(), video_playtime.count());
    video_starttime = std::chrono::high_resolution_clock::now();
#endif
    
    ++emit_seq_[stream->label_];
    src_idx = (src_idx + 1) % srcs_.size();
    stream = &srcs_[src_idx];
  }

  printf("Channel %d exit\n", channel_id_);
}

Sampler::Sampler(Config config): config_(config) {
  size_t src_idx = 0;
  for (auto src: config_["srcs"]) {
    src["src_idx"] = src_idx++;
    auto const [_, succ] = srcs_config_.try_emplace(src["id"].as<std::string>(), src);
    if (!succ) {
      throw std::runtime_error(cv::format("Duplicate src id: %s", src["id"].as<std::string>().c_str()));
    }
  }
}

void Sampler::start() {
  auto const ch_num = config_["channels"].size();
  ch_workers_.reserve(ch_num);
  windows_.resize(ch_num);

  auto const default_window_size = yaml_get_or_default(config_, "window_size", 0);

  auto const decode_thread_num = yaml_get_or_default(config_, "decode_thread_num", 1);
  decoders_.reserve(decode_thread_num);
  for (int i = 0; i < decode_thread_num; i++) {
    decoders_.emplace_back(i, config_["srcs"], &decoders_stopped_, &req_ch_);
  }

  for (size_t i = 0; i < ch_num; i++) {
    auto ch_cfg = config_["channels"][i];
    windows_[i].set_max_size(yaml_get_or_default(ch_cfg, "window_size", default_window_size));

    ch_workers_.emplace_back([this, i, ch_cfg]() {
      auto worker = ChannelWorker(i, ch_cfg, srcs_config_, windows_[i], label_emit_cnt_);
      worker.start();
    });
  }
}

void Sampler::wait() {
  for (auto& worker: ch_workers_) {
    if (worker.joinable()) worker.join();
  }
}

int Sampler::get_pkts(int ch_id, size_t num, std::vector<FramePkts>& buf, std::atomic_bool& stopped) {
  int cnt = 0;
  auto& window = windows_[ch_id];
  for (size_t i = 0; i < num; i++) {
    auto frame_opt = window.block_pop_front_unless(stopped);
    if (!frame_opt.has_value()) break;
    buf.emplace_back(std::move(frame_opt.value()));
    cnt++;
  }
  return cnt;
}

std::future<std::vector<Frame>> Sampler::decode_keyframe(FramePkts const& pkts) {
  DecodeRequest req(&pkts, DecodeRequest::ReqType::KeyFrame);
  auto ft = req.res_promise.get_future();
  
  req_ch_.emplace_back(std::move(req));
  return ft;
}

void Sampler::stop_decoders() {
  decoders_stopped_.store(true, std::memory_order_release);
  req_ch_.producer_notify_all();
  for (auto& decoder: decoders_) {
    if (decoder.worker_th_.joinable()) {
      decoder.worker_th_.join();
    }
  }
}

void Sampler::notify_all_consumers() {
  for (auto& win: windows_) {
    win.producer_notify_all();
  }
}
